# 导入日志模块
import logging
import time
import jwt
import uuid
from tools import tools
from datetime import datetime, timezone

# 从http.server模块导入HTTPServer和BaseHTTPRequestHandler类
from http.server import HTTPServer, BaseHTTPRequestHandler

# 导入json模块
import json

# 配置日志记录级别为INFO
logging.basicConfig(level=logging.INFO)

# 获取当前模块的日志记录器
logger = logging.getLogger(__name__)


class MCPServer:
    def __init__(self):
        self.protocol_version = "2025-06-18"
        self.server_info = {"name": "MCP服务器", "version": "1.0.0"}
        self.capabilities = {"tools": {}, "resources": "", "prompts": {}}
        self.sessions = {}
        # OAuth 2.1授权相关配置
        self.oauth_config = {
            # 授权端点
            "authorization_endpoint": "http://localhost:8001/oauth/authorize",
            # 令牌端点
            "token_endpoint": "http://localhost:8001/oauth/token",
            # 注册端点
            "registration_endpoint": "http://localhost:8001/oauth/register",
            # 发行者 授权服务器
            "issuer": "http://localhost:8001",
            # 支持的作用域
            "scopes_supported": ["mcp:tools", "mcp:resources", "mcp:prompts"],
            # 支持的响应类型
            "response_types_supported": ["code"],
            # 支持的授权类型
            "grant_types_supported": ["authorization_code", "refresh_token"],
            # 令牌端点支持的认证方法
            "token_endpoint_auth_methods_supported": [
                "client_secret_basic",
                "client_secret_post",
            ],
            # 支持的PKCE方法
            "code_challenge_methods_supported": ["S256"],  # PKCE支持
        }

    # 生成唯一的会话ID
    def _generate_session_id(self) -> str:
        """生成唯一的会话ID"""
        # 使用uuid4生成唯一ID
        return str(uuid.uuid4())

    def _get_or_create_session(self, session_id: str = None):
        """获取或创建会话"""
        # 如果没有传入会话ID，则生成一个新的
        if not session_id:
            session_id = self._generate_session_id()
        # 如果会话ID不存在于sessions字典中，则新建一个会话
        if session_id not in self.sessions:
            self.sessions[session_id] = {
                "created_at": time.time(),
                "last_activity": time.time(),
                "message_count": 0,
                "client_info": None,
            }
            # 记录新会话的创建
            logger.info(f"创建新会话: {session_id}")
        # 更新会话的最后活动时间
        self.sessions[session_id]["last_activity"] = time.time()
        # 增加消息计数
        self.sessions[session_id]["message_count"] += 1
        # 返回会话ID
        return session_id

    def handle_initialize(self, params, session_id):
        client_info = params.get("client_info", {})
        session_id = self._get_or_create_session(session_id)
        self.sessions[session_id]["client_info"] = client_info
        return {
            "protocolVersion": self.protocol_version,  # 协议版本
            "capabilities": self.capabilities,  # 能力
            "serverInfo": self.server_info,  # 服务器信息
            "session_id": session_id,  # 返回会话ID给客户端
        }

    def get_resource_metadata(self):
        # 返回资源元数据字典
        return {
            "resource": "http://localhost:8000",
            "authorization_servers": [
                {
                    "issuer": self.oauth_config[
                        "issuer"
                    ],  # 发行者 http://localhost:8001
                    "authorization_endpoint": self.oauth_config[
                        "authorization_endpoint"
                    ],  # 授权端点 http://localhost:8001/oauth/authorize 获取授权码的地址
                    "token_endpoint": self.oauth_config[
                        "token_endpoint"
                    ],  # 令牌端点 客户端使用授权码去哪里换令牌AccessToken
                    "registration_endpoint": self.oauth_config[
                        "registration_endpoint"
                    ],  # 注册端点 客户端去哪里自动注册
                    "scopes_supported": self.oauth_config[
                        "scopes_supported"
                    ],  # 支持的作用域  ["mcp:tools", "mcp:resources", "mcp:prompts"]
                    "response_types_supported": self.oauth_config[
                        "response_types_supported"
                    ],  # 支持的响应类型
                    "grant_types_supported": self.oauth_config[
                        "grant_types_supported"
                    ],  # 支持的授权类型
                    "token_endpoint_auth_methods_supported": self.oauth_config[
                        "token_endpoint_auth_methods_supported"
                    ],  # 令牌端点支持的认证方法  client_secret_basic 客户端注册成功后，服务器给颁发client_id client_secret
                    "code_challenge_methods_supported": self.oauth_config[
                        "code_challenge_methods_supported"
                    ],  # 支持的PKCE方法
                }
            ],
        }


class MCPHTTPHandler(BaseHTTPRequestHandler):
    def __init__(self, *args, mcp_server=None, **kwargs):
        self.mcp_server = mcp_server
        super().__init__(*args, **kwargs)

    def _handle_resource_metadata(self):
        try:
            metadata = self.mcp_server.get_resource_metadata()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps(metadata).encode("utf-8"))
        except Exception as e:
            logger.error(f"处理资源元数据时发生错误:{str(e)}")
            self.send_error(500, "Internal Server Error")

    def do_GET(self):
        try:
            if self.path == "/.well-known/oauth-protected-resource":
                self._handle_resource_metadata()
                return
            else:
                self.send_response(200)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(b"welcome to MCP Server")
                return
        except Exception as e:
            logger.error(f"GET处理错误 : {str(e)}")
            self.send_error(500, "Internal Server Error")

    def _send_unauthorized_response(self):
        self.send_response(401)
        self.send_header(
            "WWW-Authenticate", 'Bearer realm="MCP Server",error="Invalid token"'
        )
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        error_response = {
            "error": "Invalid token",
            "error_description": "the access token is invalid or expired",
        }
        self.wfile.write(json.dumps(error_response).encode("utf-8"))

    def _validate_authorization(self):
        auth_header = self.headers.get("Authorization")
        if not auth_header or not auth_header.startswith("Bearer "):
            logger.warning("缺少Authorization头或格式不正确")
            self._send_unauthorized_response()
            return False
        return True

    def _handle_mcp_message(self, message, session_id):
        msg_id = message.get("id")
        method = message.get("method")
        params = message.get("params")
        if method == "initialize":
            result = self.mcp_server.handle_initialize(params, session_id)
            return {"id": msg_id, "result": result} if msg_id else None
        elif method == "tools/list":
            # result = self.mcp_server.handle_tools_list(params, session_id)
            result = {"method": method}
            return {"id": msg_id, "result": result} if msg_id else None
        else:
            if msg_id:
                return {
                    "id": msg_id,
                    "error": {"code": 32601, "message": f"Method {method} not found "},
                }
            else:
                return None

    def do_POST(self):
        try:
            if self.path != "/mcp":
                self.send_error(404, "MCP endpoint not found")
                return
            accept_header = self.headers.get("Accept", "")
            if (
                "application/json" not in accept_header
                and "text/event-stream" not in accept_header
            ):
                self.send_error(404, "missing required accept header")
                return
            if self.path == "/mcp":
                content_length = int(self.headers.get("Content-Length", 0))
                body = self.rfile.read(content_length)
                message = json.loads(body.decode("utf-8"))
                # 如果不是初始化连接的请求，则需要验证权限，如果权限验证不通过，则直接返回
                if message and message.get("method") != "initialize":
                    if not self._validate_authorization():
                        return
            protocol_version = self.headers.get("MCP-Protocol-Version")
            if (
                protocol_version
                and protocol_version != self.mcp_server.protocol_version
            ):
                self.send_error(400, "Unsupported protocol version:{protocol_version}")
                return
            session_id = self.headers.get("MCP-Session-Id")
            print("=1=session_id", session_id)
            response = self._handle_mcp_message(message, session_id)
            if response:
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header(
                    "MCP-Protocol-Version", self.mcp_server.protocol_version
                )
                if "result" in response and "session_id" in response["result"]:
                    self.send_header("MCP-Session-Id", response["result"]["session_id"])
                self.end_headers()
                self.wfile.write(json.dumps(response).encode("utf-8"))
            else:
                self.send_response(202)
                self.end_headers()

        except Exception as e:
            logger.error(f"POST处理错误 : {str(e)}")
            self.send_error(500, "Internal Server Error")


def run_http_server(mcp_server, host, port):
    def handler_factory(*args, **kwargs):
        return MCPHTTPHandler(*args, mcp_server=mcp_server, **kwargs)

    server = HTTPServer((host, port), handler_factory)
    try:
        logger.info("MCP资源服务器成功启动")
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("MCP资源服务器已经停止运行")
    finally:
        server.server_close()


def main():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="localhost", help="主机名")
    parser.add_argument("--port", default=8000, help="端口号")
    args = parser.parse_args()
    mcp_server = MCPServer()
    run_http_server(mcp_server, args.host, args.port)


if __name__ == "__main__":
    main()
